import matplotlib.pyplot as plt
import networkx as nx
import numpy as np
import matplotlib.animation as animation


def kruskal_mst(graph):
    num_vertices = len(graph)
    edges = []

    # Collect all edges from the adjacency matrix
    for i in range(num_vertices):
        for j in range(i + 1, num_vertices):
            if graph[i][j] != 0:
                edges.append((graph[i][j], i, j))

    # Sort edges based on weight
    edges.sort()

    # Initialize disjoint set
    parent = list(range(num_vertices))
    rank = [0] * num_vertices

    def find(u):
        if parent[u] != u:
            parent[u] = find(parent[u])
        return parent[u]

    def union(u, v):
        root_u = find(u)
        root_v = find(v)
        if root_u != root_v:
            if rank[root_u] > rank[root_v]:
                parent[root_v] = root_u
            elif rank[root_u] < rank[root_v]:
                parent[root_u] = root_v
            else:
                parent[root_v] = root_u
                rank[root_u] += 1

    # Kruskal's algorithm
    mst_edges = []
    for weight, u, v in edges:
        if find(u) != find(v):
            union(u, v)
            mst_edges.append((u, v))
            yield mst_edges

            # If MST is complete (contains n-1 edges)
            if len(mst_edges) == num_vertices - 1:
                break


def animate_kruskal(graph):
    num_vertices = len(graph)
    G = nx.Graph()
    G.add_nodes_from(range(num_vertices))

    fig, ax = plt.subplots()
    pos = nx.circular_layout(G)
    nx.draw(G, pos, with_labels=True)

    def update(mst_edges):
        ax.clear()
        nx.draw(G, pos, with_labels=True)

        mst_graph = nx.Graph()
        for u, v in mst_edges:
            mst_graph.add_edge(u, v)

        nx.draw_networkx_edges(mst_graph, pos, ax=ax, edge_color='r')
        plt.title("Kruskal's MST Algorithm")

    ani = animation.FuncAnimation(fig, update, frames=kruskal_mst(graph), repeat=False, interval=500)
    plt.show()


# Example adjacency matrix for demonstration
graph = [
    [0, 7, 0, 5, 0, 0, 0],
    [7, 0, 8, 9, 7, 0, 0],
    [0, 8, 0, 0, 5, 0, 0],
    [5, 9, 0, 0, 15, 6, 0],
    [0, 7, 5, 15, 0, 8, 9],
    [0, 0, 0, 6, 8, 0, 11],
    [0, 0, 0, 0, 9, 11, 0]
]

animate_kruskal(graph)
